// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <ktest/ktest.h>

#include <nebula/array/builder_nested.h>
#include <nebula/util/list_util.h>

#include <nebula/testing/builder.h>
#include <nebula/testing/ktest_util.h>

namespace nebula {


using ListAndListViewTypes =
    ::testing::Types<ListType, LargeListType, ListViewType, LargeListViewType>;

template <typename T>
class TestListUtils : public ::testing::Test {
 public:
  using TypeClass = T;
  using offset_type = typename TypeClass::offset_type;
  using BuilderType = typename TypeTraits<TypeClass>::BuilderType;

  void SetUp() override {
    value_type_ = int16();
    type_ = std::make_shared<T>(value_type_);

    std::unique_ptr<ArrayBuilder> tmp;
    ASSERT_OK(MakeBuilder(pool_, type_, &tmp));
    builder_.reset(turbo::checked_cast<BuilderType*>(tmp.release()));
  }

  void TestRangeOfValuesUsed() {
    std::shared_ptr<ArrayData> result;

    // These list-views are built manually with the list-view builders instead
    // of using something like assert_array_from_json() because we want to test the
    // RangeOfValuesUsed() function's ability to handle arrays containing
    // overlapping list-views.

    // Empty list-like array
    ASSERT_OK(builder_->finish_internal(&result));
    builder_->reset();
    ASSERT_OK_AND_ASSIGN(auto range, list_util::internal::RangeOfValuesUsed(*result));
    ASSERT_EQ(range.first, 0);
    ASSERT_EQ(range.second, 0);

    // List-like array with only nulls
    ASSERT_OK(builder_->append_nulls(3));
    ASSERT_OK(builder_->finish_internal(&result));
    builder_->reset();
    ASSERT_OK_AND_ASSIGN(range, list_util::internal::RangeOfValuesUsed(*result));
    ASSERT_EQ(range.first, 0);
    ASSERT_EQ(range.second, 0);

    // Array with nulls and non-nulls (starting at a non-zero offset)
    Int16Builder* vb = turbo::checked_cast<Int16Builder*>(builder_->value_builder());
    ASSERT_OK(vb->append(-2));
    ASSERT_OK(vb->append(-1));
    ASSERT_OK(builder_->append(/*is_valid=*/false, 0));
    ASSERT_OK(builder_->append(/*is_valid=*/true, 2));
    ASSERT_OK(vb->append(0));
    ASSERT_OK(vb->append(1));
    ASSERT_OK(builder_->append(/*is_valid=*/true, 3));
    ASSERT_OK(vb->append(2));
    ASSERT_OK(vb->append(3));
    ASSERT_OK(vb->append(4));
    if constexpr (is_list_view_type<TypeClass>::value) {
      ASSERT_OK(vb->append(10));
      ASSERT_OK(vb->append(11));
    }
    std::shared_ptr<Array> array;
    ASSERT_OK(builder_->finish(&array));
    builder_->reset();
    ASSERT_OK(array->validate_full());
    ASSERT_OK_AND_ASSIGN(range, list_util::internal::RangeOfValuesUsed(*array->data()));
    ASSERT_EQ(range.first, 2);
    ASSERT_EQ(range.second, 5);

    // Overlapping list-views
    vb = turbo::checked_cast<Int16Builder*>(builder_->value_builder());
    ASSERT_OK(vb->append(-2));
    ASSERT_OK(vb->append(-1));
    ASSERT_OK(builder_->append(/*is_valid=*/false, 0));
    if constexpr (is_list_view_type<TypeClass>::value) {
      ASSERT_OK(builder_->append(/*is_valid=*/true, 6));
      ASSERT_OK(vb->append(0));
      ASSERT_OK(builder_->append(/*is_valid=*/true, 2));
      ASSERT_OK(vb->append(1));
      ASSERT_OK(vb->append(2));
      ASSERT_OK(vb->append(3));
      ASSERT_OK(builder_->append(/*is_valid=*/false, 0));
      ASSERT_OK(builder_->append(/*is_valid=*/true, 1));
      ASSERT_OK(vb->append(4));
      ASSERT_OK(vb->append(5));
      // -- used range ends here --
      ASSERT_OK(vb->append(10));
      ASSERT_OK(vb->append(11));
    } else {
      ASSERT_OK(builder_->append(/*is_valid=*/true, 6));
      ASSERT_OK(vb->append(0));
      ASSERT_OK(vb->append(1));
      ASSERT_OK(vb->append(2));
      ASSERT_OK(vb->append(3));
      ASSERT_OK(vb->append(4));
      ASSERT_OK(vb->append(5));
      ASSERT_OK(builder_->append(/*is_valid=*/true, 2));
      ASSERT_OK(vb->append(1));
      ASSERT_OK(vb->append(2));
      ASSERT_OK(builder_->append(/*is_valid=*/false, 0));
      ASSERT_OK(builder_->append(/*is_valid=*/true, 1));
      ASSERT_OK(vb->append(4));
    }
    ASSERT_OK(builder_->append_nulls(2));
    ASSERT_OK(builder_->finish(&array));
    builder_->reset();
    ASSERT_OK(array->validate_full());
    ASSERT_ARRAYS_EQUAL(
        *array, *assert_array_from_json(
                    type_, "[null, [0, 1, 2, 3, 4, 5], [1, 2], null, [4], null, null]"));
    // Check the range
    ASSERT_OK_AND_ASSIGN(range, list_util::internal::RangeOfValuesUsed(*array->data()));
    ASSERT_EQ(range.first, 2);
    if constexpr (is_list_view_type<TypeClass>::value) {
      ASSERT_EQ(range.second, 6);
    } else {
      ASSERT_EQ(range.second, 9);
    }
    // Check the sum of logical sizes as well
    ASSERT_OK_AND_ASSIGN(int64_t sum_of_logical_sizes,
                         list_util::internal::SumOfLogicalListSizes(*array->data()));
    ASSERT_EQ(sum_of_logical_sizes, 9);
  }

 protected:
  MemoryPool* pool_ = default_memory_pool();
  std::shared_ptr<DataType> type_;
  std::shared_ptr<DataType> value_type_;
  std::shared_ptr<BuilderType> builder_;
};

TYPED_TEST_SUITE(TestListUtils, ListAndListViewTypes);

TYPED_TEST(TestListUtils, RangeOfValuesUsed) { this->TestRangeOfValuesUsed(); }

}  // namespace nebula
